import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from matplotlib.colors import colorConverter, ListedColormap
from sklearn.utils import shuffle
from sklearn.preprocessing import OneHotEncoder


def one_hot(data, start, end):
    ohe = OneHotEncoder()
    a = np.linspace(start, end - 1, end - start)
    b = np.reshape(a, [-1, 1]).astype(np.int32)
    ohe.fit(b)
    c = ohe.transform(data).toarray()
    return c


def generate(sample_size, num_classes, diff, regression):
    np.random.seed(10)
    mean = np.random.randn(2)
    cov = np.eye(2)

    sample_per_class = int(sample_size / num_classes)

    x0 = np.random.multivariate_normal(mean, cov, sample_per_class)
    y0 = np.zeros(sample_per_class)

    for ci, d in enumerate(diff):
        x1 = np.random.multivariate_normal(mean + d, cov, sample_per_class)
        y1 = (ci + 1) * np.ones(sample_per_class)

        x0 = np.concatenate((x0, x1))
        y0 = np.concatenate((y0, y1))

    if not regression:
        y0 = np.reshape(y0, [-1, 1])
        y0 = one_hot(y0.astype(np.int32), 0, num_classes)

    x, y = shuffle(x0, y0)

    return x, y


np.random.seed(10)

num_classes = 3
x, y = generate(2000, num_classes, [[3.], [3., 0.]], False)
aa = [np.argmax(num) for num in y]
colors = ['r' if num == 0 else 'b' if num == 1 else 'y' for num in aa[:]]

plt.scatter(x[:, 0], x[:, 1], c=colors)
plt.xlabel('Scaled age (in years)')
plt.ylabel('Tumor size (in cm)')
plt.show()

input_dim = 2
lab_dim = num_classes

input_features = tf.placeholder(tf.float64, [None, input_dim])
input_labels = tf.placeholder(tf.float64, [None, lab_dim])
# 定义学习参数
weights = tf.Variable(tf.random_normal([input_dim, lab_dim], dtype=tf.float64), name='weights')
biases = tf.Variable(np.zeros([lab_dim]), name='biases')

output = tf.matmul(input_features, weights) + biases
z = tf.nn.softmax(output)

a1 = tf.argmax(z, axis=1)
b1 = tf.argmax(input_labels, axis=1)
err = tf.count_nonzero(a1 - b1)

cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=output, labels=input_labels)
loss = tf.reduce_mean(cross_entropy)

optimizer = tf.train.AdamOptimizer(0.04)
train = optimizer.minimize(loss)

max_epoch = 50
mini_batch_size = 25

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    batches = np.int32(len(y) / mini_batch_size)

    for epoch in range(max_epoch):
        sum_err = 0
        for i in range(batches):
            x1 = x[i * mini_batch_size:(i + 1) * mini_batch_size, :]
            y1 = y[i * mini_batch_size:(i + 1) * mini_batch_size, :]

            _, loss_val, output_val, err_val = sess.run([train, loss, output, err],
                                                        feed_dict={input_features: x1, input_labels: y1})
            sum_err = sum_err + (err_val / mini_batch_size)
        print('Epoch: ', '%04d' % (epoch + 1), 'cost: ', '{:.9f}'.format(loss_val),
              'err: ', sum_err / batches)

    train_x, train_y = generate(200, num_classes, [[3.], [3., 0.]], False)
    aa = [np.argmax(num) for num in train_y]
    colors = ['r' if num == 0 else 'b' if num == 1 else 'y' for num in aa[:]]

    plt.scatter(train_x[:, 0], train_x[:, 1], c=colors)

    x = np.linspace(-1, 8, 200)

    y = -x * (sess.run(weights)[0][0] / sess.run(weights)[1][0]) - sess.run(biases)[0] / sess.run(weights)[1][0]
    plt.plot(x, y, label='first line', lw=3)

    y = -x * (sess.run(weights)[0][1] / sess.run(weights)[1][1]) - sess.run(biases)[1] / sess.run(weights)[1][1]
    plt.plot(x, y, label='second line', lw=2)

    y = -x * (sess.run(weights)[0][2] / sess.run(weights)[1][2]) - sess.run(biases)[2] / sess.run(weights)[1][2]
    plt.plot(x, y, label='third line', lw=1)

    plt.legend()
    plt.show()
    print(sess.run(weights), sess.run(biases))

    train_x, train_y = generate(200, num_classes, [[3.], [3., 0.]], False)
    aa = [np.argmax(num) for num in train_y]
    colors = ['r' if num == 0 else 'b' if num == 1 else 'y' for num in aa[:]]

    plt.scatter(train_x[:, 0], train_x[:, 1], c=colors)

    x_nums = 200
    x1 = np.linspace(-1, 8, num=x_nums)
    x2 = np.linspace(-1, 8, num=x_nums)

    xx, yy = np.meshgrid(x1, x2)
    classification_plane = np.zeros((x_nums, x_nums))
    for i in range(x_nums):
        for j in range(x_nums):
            classification_plane[i, j] = sess.run(a1, feed_dict={input_features: [[xx[i, j], yy[i, j]]]})

    cmap = ListedColormap([
        colorConverter.to_rgba('r', alpha=0.30),
        colorConverter.to_rgba('b', alpha=0.30),
        colorConverter.to_rgba('y', alpha=0.30)])
    plt.contourf(xx, yy, classification_plane, cmap=cmap)
    plt.show()





